import pandas as pd
import numpy as np
from pyecharts.charts import Bar
from pyecharts import options as opts
import requests
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/90.0.4430.85 Safari/537.36"}
url = 'https://xueqiu.com/service/v5/stock/screener/quote/list?page=1&size=3679&order=desc&order_by=amount&exchange=CN&market=CN&type=sha&_=1641730868838'
b = requests.get(url=url, headers=headers)
b_data = b.json()
c = b_data['data']['list']
title = ['股票代码', '股票名称', '当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率', '市值']
all_data = []
for data in c:
    insert = [data['symbol'], data['name'], data['current'], data['chg'],
              data['percent'], data['current_year_percent'], data['volume'],
              data['amount'], data['turnover_rate'], data['pe_ttm'],
              data['dividend_yield'], data['market_capital']]
    all_data.append(insert)
output_csv = pd.DataFrame(all_data, columns=title)
output_csv.to_csv('./在上海证券交易所上市的股票.csv', index=None)

d = pd.read_csv('在上海证券交易所上市的股票.csv', dtype={'code': 'str'})
d['市值'] = d['市值'].astype(float, errors='raise')
t = ['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率', '市值']
d = d.dropna()
a = input('欢迎搜索在上海证券交易所上市的股票的数据''\n''想要专门搜索一只股票请输入1；''想要查看涨跌幅排行的请输入2；想要看市值排行榜的请输入3；')
d.set_index('股票代码', inplace=True)
while True:
    try:
        if a == '1':
            q = input('请输入你想要搜索的股票代码')
            print(d.loc[q])
            e = input('是否想要图表,y或者n')
            if e == 'y':
                da1 = d.loc[q]
                c = Bar()
                c.add_xaxis(['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率', '市值'])
                c.add_yaxis('图表',
                            da1[['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率', '市值']].tolist())
                c.set_global_opts(legend_opts=opts.LegendOpts(type_='scroll', pos_top='right', orient='vertical'),
                                  title_opts=opts.TitleOpts(title="{}".format(da1['股票名称'])),
                                  datazoom_opts=opts.DataZoomOpts(range_start=10, range_end=30))
                c.render("{}.html".format(da1['股票名称']))
                print('数据可视化结果完成,请在该目录下查找打开{}html 文件!'.format(da1['股票名字']))
        elif a == '2':
            db1 = d.sort_values('涨跌幅', ascending=False)
            db2 = db1[['股票名称', '涨跌幅']]
            print('请输入想要查询排行前几的股票')
            m = int(input())
            db3 = db2.iloc[:m]
            print(db3[['股票名称', '涨跌幅']])
            e = input('是否想要详细数据图表,y或者n')
            if e == 'y':
                bar = Bar(init_opts=opts.InitOpts(width="1700px", height="750px"))
                bar.add_xaxis(['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率', '市值'])
                for i in range(0, m):
                    db4 = db1.iloc[i]
                    bar.add_yaxis("{}".format(db4['股票名称']),
                                  db4[['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率',
                                       '市值']].tolist())
                bar.set_global_opts(
                    title_opts=opts.TitleOpts(title="涨跌幅图表"),
                    datazoom_opts=opts.DataZoomOpts(range_start=10, range_end=20),
                    legend_opts=opts.LegendOpts(type_='scroll', pos_left='right', pos_top='middle', orient='vertical'))
                bar.render("涨跌幅.html")
                print('数据可视化结果完成,请在该目录下查找打开 涨跌幅.html 文件!')
        elif a == '3':
            dc1 = d.sort_values('市值', ascending=False)
            dc2 = dc1[['股票名称', '市值']]
            print('请输入想要查询排行前几的股票')
            m = int(input())
            dc3 = dc2.iloc[:m]
            print(dc3[['股票名称', '市值']])
            e = input('是否想要详细数据图表,y或者n')
            if e == 'y':
                bar = Bar(init_opts=opts.InitOpts(width="1700px", height="750px"))
                bar.add_xaxis(['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率', '市值'])
                for i in range(0, m):
                    dc4 = dc1.iloc[i]
                    bar.add_yaxis("{}".format(dc4['股票名称']),
                                  dc4[['当前价', '涨跌额', '涨跌幅', '年初至今', '成交量', '成交额', '换手率', '市盈率(TTM)', '股息率',
                                       '市值']].tolist())
                    bar.set_global_opts(
                        title_opts=opts.TitleOpts(title="市值图表"),
                        datazoom_opts=opts.DataZoomOpts(range_start=10, range_end=20),
                        legend_opts=opts.LegendOpts(type_='scroll', pos_left='right', pos_top='middle',
                                                    orient='vertical'))
                bar.render("市值.html")
                print('数据可视化结果完成,请在该目录下查找打开 市值.html 文件!')
        else:
            print("感谢你的使用")
            break
    except KeyError:
        print('你输入的股票代码错误')
    a = input('是否继续查询，y或者n')
    if a == 'y':
        a = input('欢迎搜索在上海证券交易所上市的股票的数据想要专门搜索一只股票请输入1；''想要查看涨跌幅排行的请输入2；想要看市值排行榜的请输入3')

er = pd.read_csv('在上海证券交易所上市的股票.csv', dtype={'code': 'str'})
er = er.dropna()
df = er.iloc[:, [1, 4, 8]]
print(df)
data = er.iloc[:, [4, 8]]
data.info()

da_notnull = data
n_clusters = 3
kmean = KMeans(n_clusters)
kmean.fit(da_notnull)  
labels = kmean.labels_
kmean.predict(da_notnull)
y_ = kmean.fit_predict(da_notnull)
print(y_)
# 中心点
centers = kmean.cluster_centers_
print(centers)
# 分组
o = open('分类结果.txt', 'w+', encoding='utf-8')
for i in range(0, 3):
    content = ('第{}组'.format(i) + er[y_ == i]['股票名称'] + '\n')
    o.writelines(content)
o.close()
# 画图
plt.figure()
x = da_notnull['涨跌幅']
y = da_notnull['换手率']
plt.scatter(centers[:, 0], centers[:, 1], color='red', marker="^", s=500)
plt.scatter(da_notnull[y_ == 0].iloc[:, 0], da_notnull[y_ == 0].iloc[:, 1], color='#e24fff',
            s=5)  # 取特征第0,1列，绘制类别为1的
plt.scatter(da_notnull[y_ == 1].iloc[:, 0], da_notnull[y_ == 1].iloc[:, 1], color='blue',
            s=5)  # 取特征第0,1列，绘制类别为1的，颜色blue
plt.scatter(da_notnull[y_ == 2].iloc[:, 0], da_notnull[y_ == 2].iloc[:, 1], color='green',
            s=5)  # 取特征第0,1列，绘制类别为2的，颜色green
plt.savefig('./分类结果.png')
plt.show()